import json
from typing import Optional

from langchain_core.runnables import RunnableConfig
from pydantic import BaseModel

from inscode_agent.graph.nodes.base import BaseNodes, tool_node
from inscode_agent.mcps import MCPClientManager
from inscode_agent.models.graph import GraphState
from loguru import logger


class UseMCPToolSchema(BaseModel):
    server_name: str = None
    tool_name: str = None
    arguments: dict = {}


class MCPNodes(BaseNodes):
    def __init__(self, mcp_client_manager: MCPClientManager, verbose=False, **kwargs):
        super().__init__(verbose=verbose, **kwargs)
        self.mcp_client_manager = mcp_client_manager

    @tool_node(args_schema=UseMCPToolSchema)
    async def use_mcp_tool(self, state: GraphState, *, config: Optional[RunnableConfig] = None):
        """
        Description: Request to use a tool provided by a connected MCP server. Each MCP server can provide multiple tools with different capabilities. Tools have defined input schemas that specify required and optional parameters.
        Parameters:
        - server_name: (required) The name of the MCP server providing the tool
        - tool_name: (required) The name of the tool to execute
        - arguments: (required) A JSON object containing the tool's input parameters, following the tool's input schema
        Usage:
        <use_mcp_tool>
            <server_name>server name here</server_name>
            <tool_name>tool name here</tool_name>
            <arguments>
                {
                  "param1": "value1",
                  "param2": "value2"
                }
            </arguments>
        </use_mcp_tool>
        """
        server_name = state.goto.params.get("server_name")
        tool_name = state.goto.params.get("tool_name")
        arguments = state.goto.params.get("arguments", "{}")
        logger.info(f"execute use_mcp_tool server_name: {server_name}, tool_name: {tool_name}, arguments: {arguments}")
        if not server_name or not tool_name:
            msg = "use_mcp_tool requires server name, tool name"
            logger.error(msg)
            return self.node_failure(msg)
        try:
            if isinstance(arguments, str):
                arguments = json.loads(arguments)
            result = await self.mcp_client_manager.execute_tool(server_name, tool_name, arguments)
            content = result.content[0]
            if result.isError:
                msg = f"Error executing use_mcp_tool: {content.text}"
                return self.node_failure(msg)
            return self.node_success(f"use_mcp_tool Result: {content.text}")
        except Exception as e:
            msg = f"Error while executing use_mcp_tool: {str(e)}"
            return self.node_failure(msg)
